package com.roy.rag;

import com.roy.util.AiUtils;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.query.Metadata;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

/**
 * @author 山顶洞人郑某
 * @className VectorDemo.java 向量检索 基于Redis存储向量信息和查询向量
 * @date 2025 04 06
 */
public class VectorDemo {

    public static void main(String[] args) {
        // 创建向量模型
        OpenAiEmbeddingModel openAiEmbeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(AiUtils.API_KEY)
                .modelName(AiUtils.MODEL_NAME)
                .build();

        //基于Redis的向量存储
        RedisEmbeddingStore redisEmbeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536) //维度
                .build();
        //添加文本向量
        TextSegment textSegment1 = TextSegment.textSegment("客服电话是400-8558558");
        TextSegment textSegment2 = TextSegment.textSegment("客服工作时间是周一到周五");
        TextSegment textSegment3 = TextSegment.textSegment("客服基本信息是张三");
        TextSegment textSegment4 = TextSegment.textSegment("本店给你推荐的内容有以下信息");

        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment1).content());
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment2).content());
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment3).content());
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment4).content());

        List<EmbeddingMatch<TextSegment>> matchList = redisEmbeddingStore.findRelevant(openAiEmbeddingModel.embed("客服电话是多少").content(), 3, 0.5);
        for (EmbeddingMatch<TextSegment> textSegmentEmbeddingMatch : matchList) {
            System.out.println(textSegmentEmbeddingMatch.embedded().text()+"  "+textSegmentEmbeddingMatch.score());
        }
    }
}
